package com.mdnote.practice.dp;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * @author Rhythm-2019
 * @date 2020/10/6
 * @description 树中距离之和
 */
public class LeetCode834 {

    private int[] dp;
    private int[] size;

    private int[] ans;
    private List<List<Integer>> graph;

    public static void main(String[] args) {
        LeetCode834 leetCode834 = new LeetCode834();
        int[][] board = {{1,0}};
//        int[][] board = {{0,1},{0,2},{2,3},{2,4},{2,5}};
        int[] res = leetCode834.sumOfDistancesInTree(2, board);
        System.out.println(Arrays.toString(res));
        System.out.println(Arrays.toString(leetCode834.size));
    }

    public int[] sumOfDistancesInTree(int N, int[][] edges) {

        this.graph = new ArrayList<>();
        this.dp = new int[N];
        this.size = new int[N];
        this.ans = new int[N];

        for (int i = 0; i < N; i++) {
            this.graph.add(new ArrayList<>());
        }
        // init graph
        for (int[] edge : edges) {
            this.graph.get(edge[0]).add(edge[1]);
            // 一开始我没有加下面这句话，它不能处理【1，0】这种期刊
            this.graph.get(edge[1]).add(edge[0]);
        }
        // 计算第一个节点
        // 为什么要第二个参数，
        dfs(0, -1);
        // 换根，把节点放到根的位置
        dfs_swap(0,-1);
        return ans;
    }

    private void dfs_swap(int startNode, int distance) {
        // 不断把节点往根节点上放
        ans[startNode] = dp[startNode];
        List<Integer> childNodes = this.graph.get(startNode);
        for (Integer childNode : childNodes) {
            if (childNode == distance) {
                continue;
            }
            // 缓存当前的值
            int cur_size = size[startNode], child_size = size[childNode];
            int cur_dp = dp[startNode], child_dp = dp[childNode];

            // 修改值
            // 现在需要将startNode和childNode交换身份
            // childNode : 从子节点变成父节点，size应该加1

            dp[startNode] = dp[startNode] - dp[childNode] - size[childNode];
            size[startNode] -= size[childNode];

            dp[childNode] = dp[childNode] + dp[startNode] + size[startNode];
            size[childNode] += size[startNode];

            // 下探

            dfs_swap(childNode, startNode);

            // 恢复
            dp[startNode] = cur_dp;
            size[startNode] = cur_size;
            dp[childNode] = child_dp;
            size[childNode] = child_size;
        }
    }

    /**
     * 最优子结构： 当前节点的路径和  =  子节点的路径和之和 + 节点数量
     * @param node
     * @param startNode
     */
    private void dfs(int startNode, int distance) {
        size[startNode] = 1;
        List<Integer> childNodes = this.graph.get(startNode);
//        size[startNode] = childNodes.size() + 1;
        for (Integer childNode : childNodes) {
            if (childNode == distance) {
                continue;
            }
            dfs(childNode, startNode);
            // dp
            this.dp[startNode] += this.dp[childNode] + size[childNode];
            // size 数据记录累计大小
            this.size[startNode] += size[childNode];
        }
    }
}
